# -*- coding: utf-8 -*-
# Description:
# Setup for the OSP nmap Server
#
# Authors:
# Jan-Oliver Wagner <Jan-Oliver.Wagner@greenbone.net>
#
# Copyright:
# Copyright (C) 2015 Greenbone Networks GmbH
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

import logging
import subprocess

import defusedxml.ElementTree as secET
import nmap
from ospd.errors import OspdError
from ospd.main import main as daemon_main
from ospd.ospd import OSPDaemon, BASE_SCANNER_PARAMS

from ospd_nmap import __version__

OSPD_DESC = """
This scanner runs the tool 'nmap' to scan the target hosts.

This tool is available for most operating systems and identifies open ports,
probes the services, operating systems and even can run more sophisticated
detection routines.

For more details about nmap see the nmap homepage:
http://nmap.org/

The current version of ospd-nmap is a very simple one, only retrieving
open tcp/udp ports and service. It is possible to choice scan techniques,
to set OS detection, and to set timing and performance options.
"""

OSPD_PARAMS = {
    'mode': {
        'type': 'integer',
        'name': 'the mode of this scanner',
        'default': '1',
        'mandatory': 1,
        'visible_for_client': True,
        'description': """
            the mode: optional value includes 1, 2, 3
            1: just used to find alive hosts according the subnet.
            2: just used to find basic info according the subnet.
        """,
    },
    "args": {
        'type': 'string',
        'name': 'the args',
        'default': '',
        'mandatory': 0,
        'visible_for_client': True,
        'description': 'the args user defined, which will be executed by nmap.',
    }
}


class OSPDnmap(OSPDaemon):
    """ Class for ospd-nmap daemon. """

    def __init__(self, **kwargs):
        """ Initializes the ospd-nmap daemon's internal data. """
        super(OSPDnmap, self).__init__(**kwargs)
        self.server_version = __version__
        self.scanner_info['name'] = 'nmap'
        self.scanner_info['version'] = ''  # achieved during self.check()
        self.scanner_info['description'] = OSPD_DESC

        # used to store all the keys of dict params.
        self.__keys = [_ for _ in BASE_SCANNER_PARAMS.keys()]

        for name, param in OSPD_PARAMS.items():
            self.add_scanner_param(name, param)
            self.__keys.append(name)

        self.__nmap = nmap.PortScanner()

    def process_scan_params(self, params: dict):
        for key in params.keys():
            if not self.__keys.__contains__(key):
                raise OspdError('an invalid scanner param key: %s, just support %s' % (key, self.__keys))
        return params

    def check(self):
        """ Checks that nmap command line tool is found and is executable. """

        try:
            result = subprocess.check_output(['nmap', '-oX', '-'],
                                             stderr=subprocess.STDOUT)
        except OSError:
            # the command is not available
            return False

        if result is None:
            return False

        tree = secET.fromstring(result)

        if tree.tag != 'nmaprun':
            return False

        version = tree.attrib.get('version')
        if version == '':
            return False

        self.scanner_info['version'] = version

        return True

    def find_all_hosts(self, hosts: str, args: str = '-sP'):
        """
        find all HOSTs keeping alive in the subnet/host hosts.
        """
        self.__nmap.scan(hosts=hosts, arguments=args)
        return self.__nmap.all_hosts()

    def find_alive_host_basic_info(self, alive_host: str, ports: str, args: str = '-sS -sV -O -v'):
        """
        find the basic info, such as mac, hostname, service version and os info etc., according the host alive specified.
        alive_host must be an ip.
        """
        if 'U' in ports:
            args = args + ' -sU'

        result = self.__nmap.scan(hosts=alive_host, ports=ports, arguments=args)
        temp_rst = result['scan'][alive_host]

        ret = {'host': alive_host}

        if temp_rst['hostnames'][0]['name'] != '':
            ret['hostname'] = temp_rst['hostnames'][0]['name']
        if "mac" in temp_rst['addresses']:
            mac = temp_rst['addresses']['mac']
            ret['mac'] = mac
            if len(temp_rst['vendor']) != 0:
                ret['vendor'] = temp_rst['vendor'][mac]
        if temp_rst['status']['state'] == 'up':
            ret['status'] = temp_rst['status']['state']
        else:
            ret['status'] = temp_rst['status']['state']
            return ret

        if 'uptime' in temp_rst:
            ret['uptime'] = temp_rst['uptime']['seconds'] + 's'
        if 'tcp' in temp_rst:
            for port in temp_rst['tcp']:
                ret['tcp-service'] = temp_rst['tcp'][port]['name']
                if temp_rst['tcp'][port]['version'] != '':
                    ret['tcp-service-version'] = temp_rst['tcp'][port]['version']
                ret['tcp-service-port'] = port
                ret['tcp-service-status'] = temp_rst['tcp'][port]['state']
                if temp_rst['tcp'][port]['product'] != "":
                    ret['tcp-service-product'] = temp_rst['tcp'][port]['product']
        if len(temp_rst['osmatch']) > 0:
            ret['os-name'] = temp_rst['osmatch'][0]['name']
            # not enumerate cpe

        return ret

    def exec_scan(self, scan_id, target):
        """ Starts the nmap scanner for scan_id scan. """

        ports = self.get_scan_ports(scan_id)
        options = self.get_scan_options(scan_id)
        logging.info("the options are: %s" % options)

        mode = int(options.get('mode'))
        args = options.get('args')

        try:
            if mode == 1:
                if args == '':
                    ret = self.find_all_hosts(target)
                else:
                    ret = self.find_all_hosts(target, args)
                if ret:
                    self.add_scan_host_detail(scan_id=scan_id, host=target, name='find_hosts', value=",".join(ret))
                self.set_scan_progress(scan_id=scan_id, progress=100)
            elif mode == 2:
                alive_hosts = self.find_all_hosts(target)
                logging.info('hosts alive in the subnet %s are %s' % (target, alive_hosts.__str__()))
                ret = []
                for i, host in enumerate(alive_hosts):
                    if args == '':
                        ret.append(self.find_alive_host_basic_info(alive_host=host, ports=ports))
                    else:
                        ret.append(self.find_alive_host_basic_info(alive_host=host, ports=ports, args=args))
                    rate = (i + 1) / len(alive_hosts)
                    progress = int(100 * rate)
                    logging.info("the current progress is %s" % progress)
                    self.set_scan_progress(scan_id=scan_id, progress=progress)
                val = ret.__str__().replace("'", '"')
                if len(ret) > 0:
                    self.add_scan_host_detail(scan_id=scan_id, host=target, name='basic_info', value=val)
            else:
                raise Exception('A non-supported mode: %s' % str(mode))
        except Exception as e:
            self.add_scan_error(
                scan_id, host=target,
                value="A problem occurred trying to execute 'nmap', traceback: %s" % repr(e)
            )
            self.add_scan_error(
                scan_id, host=target,
                value="The result of 'nmap' was empty, traceback: %s" % repr(e)
            )
            return 2
        return 1


def main():
    """ OSP nmap main function. """
    daemon_main('OSPD - nmap wrapper', OSPDnmap)


if __name__ == '__main__':
    main()
